/* 
 * Quechua - the lightweight data mining framework
 *
 * Copyright (C) 2012 Marek Denis <quechua@octogan.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "dionaeaHarvester.h"


bool exec_query(dbwork_t& w,dbresult_t& r, const char* q, bool commit) {
    
    try {
        r = w.exec(q);
        if(commit)
            w.commit();
    } catch (undefined_column_ex& ex) {
        LOG(ERROR) << "Missing column? Check your SQL query and/or database, something is wrong\n"
                   << "Query executed: " << q;
        return false;
        
    }  catch (commit_ex& ex) {
        LOG(WARN) << "Not sure whether transaction was commited.";
        return false;

    } catch(std::exception& ex) {
        LOG(ERROR) << "Unhandled and unexpected exception was caught, this should not happen"
                   << ex.what() << "\n"
                   << "Query executed: " << q;
        return false;
    }

    return true;
}

DionaeaHarvester::DionaeaHarvester(): Channel(), connection(NULL) {
};
DionaeaHarvester::~DionaeaHarvester() {
    stop();
};


/**
 * What we need here is:
 * - db hostname
 * - user
 * - pass
 * - database name
 **/
bool DionaeaHarvester::load_misc_conf(const setting_t& misc) {
    string_t raw_interval = "1D";
    try {
        misc.lookupValue("user",credentials.user);
        misc.lookupValue("password",credentials.password);
        misc.lookupValue("hostname",credentials.hostname);
        misc.lookupValue("database",credentials.database);

    } catch(setting_ex& ex) {
        LOG(WARN) << "One of the required settings was not found. "
                  << "You must specify: "
                  << "-user\n-pass\n-hostname\n-database name\n-interval";
        return false;
    }
   
    Python::calculate_date_ranges(raw_interval,range.from,range.to);
    LOG(DEBUG) << "Calculated from: " << range.from << " to: " << range.to;
    return true;
}

bool DionaeaHarvester::prepare() {

    string_t connection_credentials = "dbname=" + credentials.database + 
                                      " user=" + credentials.user +
                                      " password=" + credentials.password + 
                                      " host=" + credentials.hostname;
    try {
        connection = new dbconnection_t(connection_credentials);
    } catch(alloc_ex& ex) {
        LOG(ERROR) << "DionaeaHarvester: Cannot allocate memory for connection object";
        return false;
    } catch(connection_ex& ex) {
        LOG(WARN) << "DionaeaHarvester: Database connection problem";
        delete connection;
        connection = NULL;
        return false;
    } catch(std::exception& ex) {
        LOG(WARN) << "DionaeaHarvester: Unexpected exception catched: " << ex.what();
        delete connection;
        connection = NULL;
        return false;
    }

    LOG(INFO) << "Connected to "<< credentials.user << "@" 
              << credentials.hostname << ":" << credentials.database;
    return true;
};

bool DionaeaHarvester::notify() {
    if(connection) {
        fetch_data();
        return true;
    }
    return false;
}

bool DionaeaHarvester::start() {
    return true;
}

bool DionaeaHarvester::stop() {
    if(connection) {
        connection->disconnect();
        LOG(INFO) << "Disconnected from "<< credentials.user << "@"
                  << credentials.hostname << ":" << credentials.database;
        connection=NULL;
    }
};

void DionaeaHarvester::fetch_data() {
    static char buf[1024];
#ifdef SELECT_V1
    sprintf(buf,ChannelSelects::fetch_connections);
#else
    snprintf(buf,1023,ChannelSelects::fetch_connections,range.from.c_str(),range.to.c_str());
#endif
    buf[1023] = '\0';

    dbwork_t work(*connection);
    dbresult_t result;

    exec_query(work,result,buf,true);

    DionaeaConnections *connections = new DionaeaConnections();

    DionaeaConnections::unindexedTransaction* txn;
    for(int i=0;i<result.size();++i) {
        txn = new DionaeaConnections::unindexedTransaction();
        txn->protocol = result[i][2].as<const char*>();
        txn->dst_host = result[i][3].as<const char*>();
        txn->dst_port = result[i][4].as<u_int16_t>();
        txn->src_host = result[i][5].as<const char*>();
        txn->src_port = result[i][6].as<u_int16_t>();
        connections->rows.push_back(txn);
    }
    
    connections->set_dbresult(result);
    connections->set_date_range(range.from,range.to);

    notify_workflows(connections);
    connections = NULL; //forget about it
};

/* A SIMPLE FACTORY FOR CREATING OBJECTS */
extern "C" Channel* create_channel() {
    return new DionaeaHarvester();
};

